{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/meta-llama/llama-recipes/blob/main/recipes/3p_integrations/llamaindex/dlai_agentic_rag/Building_Agentic_RAG_with_Llamaindex_L2_Tool_Calling.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "This notebook ports the DeepLearning.AI short course [Building Agentic RAG with Llamaindex Lesson 2 Tool Calling](https://learn.deeplearning.ai/courses/building-agentic-rag-with-llamaindex/lesson/3/tool-calling) to using Llama 3. It shows how to use Llama 3 to not only pick a function to execute, but also infer an argument to pass through the function.\n",
    "\n",
    "You should take the course before or after going through this notebook to have a deeper understanding.\n",
    "\n",
    "Note: Unlike Lesson 1 where we use Llama 3 70b on [Groq](https://groq.com/), this lesson uses Llama 3 on [Fireworks.ai](https://fireworks.ai/) to overcome the rate limit issue with Groq on some summary tool calling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eiJsOa29ej7G",
    "outputId": "edc5d39c-f379-4410-db9f-998db9c099be"
   },
   "outputs": [],
   "source": [
    "!pip install llama-index\n",
    "!pip install llama-index-embeddings-huggingface\n",
    "!pip install llama-index-llms-fireworks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NZ9l6k_3WncE"
   },
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QkaALpnIQ01b"
   },
   "outputs": [],
   "source": [
    "from llama_index.core.tools import FunctionTool\n",
    "\n",
    "def add(x: int, y: int) -> int:\n",
    "    \"\"\"Adds two integers together.\"\"\"\n",
    "    return x + y\n",
    "\n",
    "def mystery(x: int, y: int) -> int:\n",
    "    \"\"\"Mystery function that operates on top of two numbers.\"\"\"\n",
    "    return (x + y) * (x + y)\n",
    "\n",
    "\n",
    "add_tool = FunctionTool.from_defaults(fn=add)\n",
    "mystery_tool = FunctionTool.from_defaults(fn=mystery)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "\n",
    "os.environ['FIREWORKS_API_KEY'] = 'xxx' # get a free key at https://fireworks.ai/api-keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mA3AG6CFQ3fj",
    "outputId": "b872d91f-3a16-4d40-cacb-59c8ba5b5bde"
   },
   "outputs": [],
   "source": [
    "from llama_index.llms.fireworks import Fireworks\n",
    "\n",
    "# Llama 3 8b on Fireworks.ai also works in some cases, but 70b works better overall\n",
    "#llm = Fireworks(model=\"accounts/fireworks/models/llama-v3-8b-instruct\", temperature=0)\n",
    "llm = Fireworks(model=\"accounts/fireworks/models/llama-v3-70b-instruct\", temperature=0)\n",
    "\n",
    "# a quick sanity test\n",
    "#llm.complete(\"Who wrote the  book godfather? \").text\n",
    "\n",
    "response = llm.predict_and_call(\n",
    "    [add_tool, mystery_tool],\n",
    "    \"Tell me the output of the mystery function on 2 and 9\",\n",
    "    verbose=True\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8TTkd6vuUMmh",
    "outputId": "c2c419af-e9d1-48bb-aa51-c785dcdee3a0"
   },
   "outputs": [],
   "source": [
    "!wget \"https://openreview.net/pdf?id=VtmBAGCN7o\" -O metagpt.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "auZQalH5J7CU"
   },
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "# https://arxiv.org/pdf/2308.00352 metagpt.pdf\n",
    "documents = SimpleDirectoryReader(input_files=[\"metagpt.pdf\"]).load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GFfUjJypJ7Eq"
   },
   "outputs": [],
   "source": [
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "splitter = SentenceSplitter(chunk_size=1024)\n",
    "nodes = splitter.get_nodes_from_documents(documents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "FrqorjH3VHmT",
    "outputId": "b4888caf-0623-4d64-dba1-4d74c12e64f3"
   },
   "outputs": [],
   "source": [
    "print(nodes[0].get_content(metadata_mode=\"all\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9WisqWK4VPCZ"
   },
   "outputs": [],
   "source": [
    "from llama_index.core import Settings, VectorStoreIndex\n",
    "from llama_index.embeddings.huggingface import HuggingFaceEmbedding\n",
    "\n",
    "Settings.llm = llm\n",
    "\n",
    "Settings.embed_model = HuggingFaceEmbedding(\n",
    "    model_name=\"BAAI/bge-small-en-v1.5\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YS4e0mzsVKsl"
   },
   "outputs": [],
   "source": [
    "# Settings.llm and embed_model apply to which call below? VectorStoreIndex(), as_query_engine?\n",
    "\n",
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "vector_index = VectorStoreIndex(nodes)\n",
    "query_engine = vector_index.as_query_engine(similarity_top_k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "S7tz2Z28VKv1"
   },
   "outputs": [],
   "source": [
    "from llama_index.core.vector_stores import MetadataFilters\n",
    "\n",
    "query_engine = vector_index.as_query_engine(\n",
    "    similarity_top_k=2,\n",
    "    filters=MetadataFilters.from_dicts(\n",
    "        [\n",
    "            {\"key\": \"page_label\", \"value\": \"2\"}\n",
    "        ]\n",
    "    )\n",
    ")\n",
    "\n",
    "response = query_engine.query(\n",
    "    \"What are some high-level results of MetaGPT?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "CttWxW8aVKyk",
    "outputId": "4b64a64f-a989-4ee0-f08e-a6d6f5db42b6"
   },
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZvQGoUR0VK1I",
    "outputId": "9033f46f-baba-4345-bd6c-29ce4db3ea39"
   },
   "outputs": [],
   "source": [
    "for n in response.source_nodes:\n",
    "    print(n.metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5r1MHbLOPT8Y"
   },
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from llama_index.core.vector_stores import FilterCondition\n",
    "\n",
    "\n",
    "def vector_query(\n",
    "    query: str,\n",
    "    page_numbers: List[str]\n",
    ") -> str:\n",
    "    \"\"\"Perform a vector search over an index.\n",
    "\n",
    "    query (str): the string query to be embedded.\n",
    "    page_numbers (List[str]): Filter by set of pages. Leave BLANK if we want to perform a vector search\n",
    "        over all pages. Otherwise, filter by the set of specified pages.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    metadata_dicts = [\n",
    "        {\"key\": \"page_label\", \"value\": p} for p in page_numbers\n",
    "    ]\n",
    "\n",
    "    query_engine = vector_index.as_query_engine(\n",
    "        similarity_top_k=2,\n",
    "        filters=MetadataFilters.from_dicts(\n",
    "            metadata_dicts,\n",
    "            condition=FilterCondition.OR\n",
    "        )\n",
    "    )\n",
    "    response = query_engine.query(query)\n",
    "    return response\n",
    "\n",
    "\n",
    "vector_query_tool = FunctionTool.from_defaults(\n",
    "    name=\"vector_tool\",\n",
    "    fn=vector_query\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "2jMB3iS6VjFg",
    "outputId": "4ecd2d26-e159-4765-9aa5-818c5308ae02"
   },
   "outputs": [],
   "source": [
    "response = llm.predict_and_call(\n",
    "    [vector_query_tool],\n",
    "    \"What are the high-level results of MetaGPT as described on page 2?\",\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "so2p09VNVm9I",
    "outputId": "8fd7027b-e356-492c-decf-36340ad90978"
   },
   "outputs": [],
   "source": [
    "for n in response.source_nodes:\n",
    "    print(n.metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AuxhlFxHV4MV"
   },
   "outputs": [],
   "source": [
    "from llama_index.core import SummaryIndex\n",
    "from llama_index.core.tools import QueryEngineTool\n",
    "\n",
    "summary_index = SummaryIndex(nodes)\n",
    "summary_query_engine = summary_index.as_query_engine(\n",
    "    response_mode=\"tree_summarize\",\n",
    "    use_async=True,\n",
    ")\n",
    "summary_tool = QueryEngineTool.from_defaults(\n",
    "    name=\"summary_tool\",\n",
    "    query_engine=summary_query_engine,\n",
    "    description=(\n",
    "        \"Useful if you want to get a summary of MetaGPT\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dSqh6iLkV61F",
    "outputId": "9dc5e545-9f9d-4a7d-8332-beb158574b8e"
   },
   "outputs": [],
   "source": [
    "response = llm.predict_and_call(\n",
    "    [vector_query_tool, summary_tool],\n",
    "    \"What are the MetaGPT comparisons with ChatDev described on page 8?\",\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OsPn62A2V8R7",
    "outputId": "cf74caa0-5f59-4f3c-f806-f996b317df8c"
   },
   "outputs": [],
   "source": [
    "for n in response.source_nodes:\n",
    "    print(n.metadata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ugdDXz1EV96J",
    "outputId": "fdc93775-04a0-4632-f190-0ef8d3800651"
   },
   "outputs": [],
   "source": [
    "response = llm.predict_and_call(\n",
    "    [vector_query_tool, summary_tool],\n",
    "    \"What is a summary of the paper?\",\n",
    "    verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fx0RBeX0OS3n"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
